import torch.nn.functional as F

# 训练损失
def model_loss_train(disp_ests, disp_gts, img_masks): # 输入预测视差List[上下文上采样视差(B,H,W),1/4原尺寸回归视差(B,H/4,W/4)]，地面实况List[全尺寸地面实况,1/4原尺寸地面实况]，掩码List[全尺寸掩码,1/4原尺寸掩码]
    weights = [1.0, 0.3] # 全尺寸权重1.0，1/4尺寸权重0.3
    all_losses = []
    for disp_est, disp_gt, weight, mask_img in zip(disp_ests, disp_gts, weights, img_masks):
        all_losses.append(weight * F.smooth_l1_loss(disp_est[mask_img], disp_gt[mask_img], reduction='mean')) # 所有损失加入：权重*smoothL1损失
    return sum(all_losses) # 所有损失求和

# 验证损失
def model_loss_test(disp_ests, disp_gts, img_masks): # 输入预测视差List[上下文上采样视差(B,H,W)]，地面实况List[全尺寸地面实况]，掩码List[全尺寸掩码]
    weights = [1.0] # 全尺寸权重1.0
    all_losses = []
    for disp_est, disp_gt, weight, mask_img in zip(disp_ests, disp_gts, weights, img_masks):
        all_losses.append(weight * F.l1_loss(disp_est[mask_img], disp_gt[mask_img], reduction='mean')) # 所有损失加入：权重*L1损失
    return sum(all_losses) # 所有损失求和
